import tensorflow as tf  # 导入 TensorFlow 库，用于构建和训练深度学习模型
from keras import layers, models  # 从 TensorFlow 的 Keras API 导入层和模型模块
import numpy as np  # 导入 NumPy 库，用于数值计算和数组操作
import os  # 导入 os 模块，用于文件路径操作
from PIL import Image  # 从 Pillow 库导入 Image 模块，用于图像处理
from sklearn.model_selection import train_test_split  # 从 scikit-learn 导入 train_test_split 函数，用于划分训练集和测试集


# 定义数据加载和预处理函数
def load_and_preprocess_data(data_dir):  # 定义一个函数，用于加载和预处理数据
    images = []  # 初始化一个空列表，用于存储图像数据
    labels = []  # 初始化一个空列表，用于存储标签数据
    label_mapping = {'X': 1, 'V': 0, 'O': 2}  # 定义标签映射，将字符映射为数字
    for root, dirs, files in os.walk(data_dir):  # 遍历指定目录及其子目录
        for file in files:  # 遍历当前目录下的所有文件
            if file.endswith(('.png', '.jpg', '.jpeg')):  # 检查文件是否为图片格式
                label = os.path.basename(root)  # 获取当前目录的名称作为标签
                if label in label_mapping:  # 检查标签是否在映射中
                    img_path = os.path.join(root, file)  # 构造图片的完整路径
                    img = Image.open(img_path).convert('L')  # 打开图片并转换为灰度图
                    img = img.resize((28, 28))  # 调整图片大小为 28x28 像素
                    img_array = np.array(img) / 255.0  # 将图片转换为 NumPy 数组并归一化
                    img_array = np.expand_dims(img_array, axis=-1)  # 增加一个维度，使形状变为 (28, 28, 1)
                    images.append(img_array)  # 将处理后的图片添加到 images 列表
                    labels.append(label_mapping[label])  # 将对应的标签添加到 labels 列表
    images = np.array(images)  # 将 images 列表转换为 NumPy 数组
    labels = np.array(labels)  # 将 labels 列表转换为 NumPy 数组
    return images, labels  # 返回处理后的图像和标签数据


# 加载数据
data_dir = '/home/ai/image2text/dataset_train_raw/'  # 替换为你的数据目录
images, labels = load_and_preprocess_data(data_dir)  # 调用函数加载和预处理数据

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(images, labels, test_size=0.2, random_state=42)  # 使用 train_test_split 划分数据集

# 构建模型
model = models.Sequential([  # 创建一个顺序模型
    layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),  # 添加卷积层，32 个滤波器，3x3 卷积核，ReLU 激活函数，输入形状为 (28, 28, 1)
    layers.MaxPooling2D((2, 2)),  # 添加最大池化层，池化窗口大小为 2x2
    layers.Conv2D(64, (3, 3), activation='relu'),  # 添加卷积层，64 个滤波器，3x3 卷积核，ReLU 激活函数
    layers.MaxPooling2D((2, 2)),  # 添加最大池化层，池化窗口大小为 2x2
    layers.Flatten(),  # 添加展平层，将二维特征图展平为一维向量
    layers.Dense(64, activation='relu'),  # 添加全连接层，64 个神经元，ReLU 激活函数
    layers.Dense(3, activation='softmax')  # 添加输出层，3 个神经元，Softmax 激活函数
])

# 编译模型
model.compile(optimizer='adam',  # 使用 Adam 优化器
              loss='sparse_categorical_crossentropy',  # 使用稀疏分类交叉熵损失函数
              metrics=['accuracy'])  # 监控准确率指标

# 训练模型
model.fit(X_train, y_train, epochs=10, validation_data=(X_test, y_test))  # 训练模型，10 个 epoch，使用验证集进行评估

# 评估模型
test_loss, test_acc = model.evaluate(X_test, y_test)  # 评估模型在测试集上的性能
print(f"Test accuracy: {test_acc}")  # 打印测试集上的准确率

# 保存模型
model.save('character_recognition_model.h5')  # 将训练好的模型保存为 H5 文件